﻿# -*- coding: utf-8 -*-
import random
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from network_env import NetworkEnvironment, draw_network
from matplotlib.animation import FuncAnimation
import pickle

# 设置中文字体为黑体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 或使用其他中文字体，如 'Microsoft YaHei'
# 解决负号显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False
   
env = NetworkEnvironment(node_count=125, link_count=806)
G = env.graph
draw_network(G)
nodes = {n: list(G.neighbors(n)) for n in G.nodes()}  # 邻接表
source = env.source
target = env.destination
max_hops = 150 # 根据网络直径自动调整更佳

# 路径生成与验证
def generate_valid_path():
    path = [source]
    current = source
    visited = set()
    
    for _ in range(max_hops):
        if current == target:
            break
        valid_neighbors = [n for n in nodes[current] if n not in visited]
        if not valid_neighbors:
            break
        next_node = random.choice(valid_neighbors)
        path.append(next_node)
        visited.add(next_node)
        current = next_node
    
    # 补齐路径长度
    while len(path) <= max_hops and path[-1] != target:
        path.append(target)
    return path[:max_hops+1]

# 适应度函数加强版（考虑路径有效性）
def enhanced_fitness(path):
    if path[0] != source or path[-1] != target:
        return float('inf')
    
    # 检查路径连通性
    for i in range(len(path)-1):
        if path[i+1] not in nodes[path[i]]:
            return float('inf')
    
    return len(path) - 1  # 跳数优先

# 增强型PSO优化器
class EnhancedPSO:
    def __init__(self, num_particles=100, max_iter=200):
        self.num_particles = num_particles
        self.max_iter = max_iter
        self.gbest = None
        self.fitness_history= []
        
        # 参数自适应调整
        self.w = 0.7
        self.c1 = 1.4
        self.c2 = 1.5
    
    def optimize(self):
        # 粒子初始化
        # swarm是一个列表，包含每个粒子的初始路径、个体最优路径和适应度
        swarm = [{
            'position': generate_valid_path(),
            'pbest': None,  # 个体最优路径
            'pbest_fit': float('inf')
        } for _ in range(self.num_particles)] 
        
        # 初始化全局最优
        self.gbest = min(swarm, 
                       key=lambda x: enhanced_fitness(x['position']))['position'].copy() 
        #这一句理解为先min(swarm,key=lambda x: enhanced_fitness(x['position']))找出符合的字典（swarm的一个元素）
        #再通过['position']取出这个字典的'position'值
        #最后执行.copy()方法，把这个值复制到self.gbest中
        # key是min函数的参数，x是swarm列表中的每个元素

        for iter in range(self.max_iter):
            # 参数衰减
            self.w *= 0.98
            self.c1 *= 0.995
            self.c2 *= 0.995
            
            for particle in swarm:
                current_fit = enhanced_fitness(particle['position'])
                
                # 更新个体最优
                if current_fit < particle['pbest_fit']:
                    particle['pbest'] = particle['position'].copy()
                    particle['pbest_fit'] = current_fit
                
                # 更新全局最优
                if current_fit < enhanced_fitness(self.gbest):
                    self.gbest = particle['position'].copy()
                
                # 速度更新（基于路径相似性）
                new_path = self._update_position(particle)
                particle['position'] = new_path
            
            # 记录收敛过程
            self.fitness_history.append(enhanced_fitness(self.gbest))
           
            print(f"\n第{iter+1}次迭代\n跳数：{self.fitness_history[-1]}")
            #print(f"Iteration {iter+1}: Best Hops: {self.fitness_history[-1]}")
            print("Optimized Path:", self.gbest)
            print(f"Total Hops: {len(self.gbest)-1}")
            
        return self.gbest
    
    def _update_position(self, particle):
        new_path = [source]
        current_node = source
        
        for step in range(1, max_hops+1):
            if current_node == target:
                break
            
            # 动态候选节点选择
            candidates = []
            neighbors = nodes[current_node]
            
            # 加入各参考方向的节点
            ref_nodes = [
                particle['position'][step] if step < len(particle['position']) else None,
                particle['pbest'][step] if particle['pbest'] and step < len(particle['pbest']) else None,
                self.gbest[step] if step < len(self.gbest) else None
            ]
            
            # 概率选择策略
            for ref in ref_nodes:
                if ref in neighbors:
                    candidates.append(ref)
            
            # 加入随机探索
            if random.random() < 0.2:
                candidates.extend(neighbors)
            
            if not candidates:
                candidates = neighbors.copy()
            
            candidates = list(set(candidates))
            
            # 概率选择（适应度导向）
            if candidates:
                next_node = self._probability_select(candidates, current_node)
                new_path.append(next_node)
                current_node = next_node
            else:
                break
        
        # 路径补齐
        while len(new_path) <= max_hops and new_path[-1] != target:
            new_path.append(target)
        return new_path[:max_hops+1]
    
    def _probability_select(self, candidates, current):
        # 基于到目标节点的距离进行概率加权
        distances = []
        for node in candidates:
            try:
                distance = nx.shortest_path_length(G, node, target)
            except:
                distance = 100
            distances.append(1 / (distance + 1))
        
        total = sum(distances)
        probs = [d/total for d in distances]
        return np.random.choice(candidates, p=probs)

def animate_smooth_path(G, path, frames_per_segment=20):
    fig, ax = plt.subplots(figsize=(12, 12))
    pos = nx.get_node_attributes(G, 'pos')  # 使用真实坐标
    
    # 绘制静态网络
    nx.draw_networkx_nodes(G, pos, node_color='lightblue', ax=ax)
    nx.draw_networkx_edges(G, pos, edge_color='gray', alpha=0.4, ax=ax)
    ax.set_title("Smooth Path Animation")


    # 生成移动轨迹数据
    trace = generate_smooth_trace(path, pos, frames_per_segment)
    
    # 初始化移动红点
    red_dot, = ax.plot([], [], 'o', color='blue', 
                      markersize=15, animated=True)
    
    # 添加路径序号标注
    node_labels = {node: f"{i}" for i, node in enumerate(path)}
    node_labels[path[0]] = f"起点\n{path[0]}"
    node_labels[path[-1]] = f"终点\n{path[-1]}"
    label_objects = nx.draw_networkx_labels(G, pos, labels=node_labels, ax=ax)
    
    def animate(i):
        # 清除上一帧的高亮边（修正部分）
        while len(ax.collections) > 2:  # 保留原始网络边
            ax.collections[-1].remove()
    
        # 高亮最优路径（静态部分）
        path_edges = list(zip(path, path[1:]))
        nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='red', ax=ax)
        nx.draw_networkx_edges(G, pos, edgelist=path_edges, 
                              edge_color='red', width=2, ax=ax)

        # 更新红点位置
        x, y = trace[i]
        red_dot.set_data([x], [y])
    
        # 动态高亮路径段
        current_segment = get_current_segment(i, len(trace), frames_per_segment)
        if current_segment is not None and (current_segment+1) < len(path):
            edge = (path[current_segment], path[current_segment+1])
            if edge[0] in G and edge[1] in G[edge[0]]:
                nx.draw_networkx_edges(
                    G, pos, edgelist=[edge], 
                    edge_color='red', width=3, 
                    ax=ax, arrows=True
                )
    
        return red_dot,

    ani = FuncAnimation(
        fig, animate,
        frames=len(trace),
        interval=50,
        blit=True,
        repeat=True
    )
    plt.title("电网无线传感网络路由可视化")
    plt.show()
    plt.close()
    return ani

def generate_smooth_trace(path, pos, frames_per_segment):  
    # 生成平滑移动轨迹坐标序列"""
    trace = [] # 存储轨迹点（每一帧）
    
    for i in range(len(path)-1):
        start = np.array(pos[path[i]])
        end = np.array(pos[path[i+1]])
        
        # 计算路径段长度
        distance = np.linalg.norm(end - start)
        # 动态计算插值点数（至少5帧）
        frames = max(2, int(distance * frames_per_segment / 10))  # 让帧数和距离成正比，为了匀速


        # 线性插值生成中间点
        for t in np.linspace(0, 1, frames):  #frames_per_segment会让每段时间一样但是不匀速
            point = start + t*(end - start)
            trace.append(point)
    # 终点停留（保持总帧数比例）？？？
    if path:
        stay_frames = int(0.1 * len(trace))  # 停留帧数=总帧数的10%
        #trace.extend([np.array(pos[path[-1]])] * stay_frames)
      
    # 在终点停留数帧
    #trace.extend([np.array(pos[path[-1]])]*5)
    return trace

def get_current_segment(frame_index, total_frames, frames_per_segment):
    #计算当前所在的路径段索引"""
    max_segment = (total_frames - 5) // frames_per_segment  # 排除最后的停留帧
    current_segment = frame_index // frames_per_segment
    return current_segment if current_segment < max_segment else None

def plot_iteration_hops(pso):

    # 数据验证
    if pso.max_iter != len(pso.fitness_history):
        raise ValueError("iterations和fitness_history长度必须相同")
    
    # 创建画布
    plt.figure()
    
    # 绘制主曲线
    plt.plot(range(1, pso.max_iter + 1), pso.fitness_history, 
             markerfacecolor='white', 
             markeredgewidth=2, 
             label='实际跳数')
    
    # 绘制理论最优线
    if pso.gbest is not None:
        plt.axhline(y=len(best_path)-1, 
                   linestyle=':', 
                   linewidth=3,
                   label='理论最优')
    
    
    # 图表装饰
    plt.title("路径优化收敛过程分析", fontsize=14, pad=20)
    plt.xlabel("迭代次数", fontsize=12)
    plt.ylabel("路径跳数", fontsize=12)
    #plt.xticks(np.arange(min(iterations), max(iterations)+1, step=max(1, len(iterations)//10)))
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.legend(loc='upper right')
    
    plt.tight_layout()
    plt.show()


# 主程序
if __name__ == "__main__":
    # 优化过程
    pso = EnhancedPSO(num_particles=50, max_iter=50) 
    best_path = pso.optimize()
    


    # 结果可视化
    print("\n\n最终结果")
    print("最优路径", best_path)
    print(f"跳数: {len(best_path)-1}")
    
    plot_iteration_hops(pso)
    ani = animate_smooth_path(G, best_path)
    ani.save("path_animation.gif", writer='pillow', fps=15)  # 保存最优路径动态图
    #ani.save("path_animation.gif", writer="imagemagick", fps=15)

    # 保存 env 对象到文件
    with open('env_instance.pkl', 'wb') as f: pickle.dump(env, f)

    #with open('env_instance.pkl', 'rb') as f: aded_env = pickle.load(f)